from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from operator import itemgetter

model = ChatOpenAI()

prompt1 = ChatPromptTemplate.from_template("what is the city {person} is from ?")
prompt2 = ChatPromptTemplate.from_template("what country is the city {city} in ? respond in {language}")

chain1 = prompt1 | model | StrOutputParser()

chain2 = (
    {"city": chain1, "language": itemgetter("language")}
    | prompt2
    | model
    | StrOutputParser()
)

result = chain2.invoke({"person": "obama", "language": "spanish"})
print(result)

